
import arrow
from datetime import datetime
from functools import wraps

from flask_jwt_extended import JWTManager
from flask_jwt_extended import create_access_token
from flask_jwt_extended import get_jwt, verify_jwt_in_request
from flask_jwt_extended import get_jwt_identity
from flask_jwt_extended import set_access_cookies
from flask_jwt_extended.config import config

from app.auth.scope import NO_NEED_JWT_ENDPOINT
from app.errors import Invalid_Token, Expired_Token, Missing_Token, PermissionDenied, NotFound
from flask import g, request, current_app
from app.db.models.users import User, Token
from app.db.models.api_token import ApiToken
from app.config.settings import JWT_ACCESS_COOKIE_NAME

jwt = JWTManager()


def refresh_expiring_jwts(response):
    try:
        exp_timestamp = get_jwt()["exp"]
        now_timestamp = datetime.timestamp(datetime.now())
        if now_timestamp > exp_timestamp:
            access_token = create_access_token(identity=get_jwt_identity())
            set_access_cookies(response, access_token)
        return response
    except (RuntimeError, KeyError):
        # Case where there is not a valid JWT. Just return the original response
        return response


def get_operator():
    # 通过参数_username传递
    if operator := request.args.get('_username'):
        return operator

    # 通过参数operator传递
    if operator := request.args.get('operator'):
        return operator

    # 通过json body operator 参数传递
    if isinstance(json := request.json, dict) and (operator := json.get('operator')):
        return operator

    return g.user.username


def g_context():
    if request.endpoint in NO_NEED_JWT_ENDPOINT or request.method in config.exempt_methods:
        return

    if request.endpoint and request.endpoint.startswith("v3"):
        token = request.cookies.get("token", "")
        token_obj = ApiToken.objects.filter(token=token, is_active=True).first()
        if not token_obj:
            raise Missing_Token()
        if arrow.get(token_obj.expired_time).timestamp() < arrow.now().timestamp():
            raise Expired_Token()
        g.identity = token_obj.uid
        g.entity_type = 'user'
    else:
        verify_jwt_in_request()
        # 验证是不是黑名单token
        token = request.cookies.get(JWT_ACCESS_COOKIE_NAME, "")
        if token and current_app.extensions["redis"].get(f"scanv_black_token_{token}"):
            raise Expired_Token()
        g.identity = get_jwt_identity()
        g.entity_type = get_jwt().get('entityType', 'user')
    g.log_type = 'ops'
    if g.entity_type == 'user':
        g.logSource = 'front'
        g.user = User.objects.with_id(g.identity)
        if not g.user:
            raise Invalid_Token()
        g.role = g.user.roleIds[0]
        g.audit_uid = g.user.pk
    elif g.entity_type == 'app':
        g.logSource = 'back'
        entity = Token.objects.with_id(g.identity)
        if not entity:
            raise Invalid_Token()
        g.role = entity.roleIds[0]
        g.user = User.objects.filter(id=entity.createBy).first()
        g.audit_uid = entity.pk
    else:
        raise Invalid_Token()

    g.operator = get_operator()
    g.ip = request.headers.get('X-Forwarded-For', request.remote_addr)


def permission_can(permission):
    """
    检测用户是否有特定权限
    :param permission
    :return:
    """
    role_map = {
        'admin': '管理员',
        'user': '普通用户'
    }
    return g.role.name == role_map[permission]


def permission_required(permission):
    """
    权限认证装饰器
    :param permission:
    :return:
    """

    def decorator(f):
        @wraps(f)
        def decorated_function(*args, **kwargs):
            if not permission_can(permission):
                raise PermissionDenied()
            return f(*args, **kwargs)

        return decorated_function

    return decorator


@jwt.expired_token_loader
def my_expired_token_callback(jwt_header, jwt_payload):
    return Expired_Token(), Expired_Token.code


@jwt.invalid_token_loader
def my_invalid_token_callback(error_string):
    return Invalid_Token(), Invalid_Token.code


@jwt.unauthorized_loader
def my_missing_token_callback(error_string):
    """
    :param error_string: String indicating why the token is invalid
    """
    return Missing_Token(), Missing_Token.code
